import torch
import time

# 设置矩阵大小
matrix_size = 10000

# 创建随机矩阵
A = torch.randn(matrix_size, matrix_size)
B = torch.randn(matrix_size, matrix_size)

print('cuda is_available: ', torch.cuda.is_available())
# 在 CPU 上运行
start_time = time.time()
C_cpu = torch.matmul(A, B)
cpu_time = time.time() - start_time
print(f"CPU 耗时: {cpu_time:.6f} 秒")

# 检查是否有可用的 GPU
if torch.cuda.is_available():
    # 将矩阵移动到 GPU
    A_gpu = A.cuda()
    B_gpu = B.cuda()

    # 在 GPU 上运行
    start_time = time.time()
    C_gpu = torch.matmul(A_gpu, B_gpu)
    gpu_time = time.time() - start_time
    print(f"GPU 耗时: {gpu_time:.6f} 秒")

    # 比较 CPU 和 GPU 的结果是否一致
    if torch.allclose(C_cpu, C_gpu.cpu(), rtol=1e-4):
        print("CPU 和 GPU 计算结果一致！")
    else:
        print("CPU 和 GPU 计算结果不一致！")
else:
    print("未检测到 GPU,跳过 GPU 测试。")
